import numpy as np
import matplotlib.pyplot as plt
from run import eval_GA
from run import eval_DFF
from network import K
import network
from DQN_FILE import DQN
# 基因长度
GENE_SIZE = network.TASK_NUM
#基因交叉长度
GENE_CROSS=1
#基因交叉概率

CROSS_RATE=0.9
#基因变异长度
GENE_MUTATION=1
#基因变异概率
MUTATION_RATE=0.01
# 种群大小
POPULATION_SIZE = 100
# 保留个体数
REMAIN_NUM = 30
# 迭代轮次
EPOCH = 3000


# 获取单个个体的适应度
def get_fitness(gene_list,env):
    res=[]
    for gene in gene_list:
        score = eval_GA(gene,env)
        res.append(score)
    return res

def select(pop, fitness):
    idx = np.random.choice(np.arange(POPULATION_SIZE), size=POPULATION_SIZE, replace=True,
                           p=fitness / (np.sum(fitness)))
    return pop[idx]

def cross(father,mother):
    idx=np.random.choice(np.arange(GENE_SIZE),size=GENE_CROSS)
    for i in range(GENE_SIZE):
        if i in idx:
            father[i]=mother[i]
def mutation(child):
    idx = np.random.choice(np.arange(GENE_SIZE), size=GENE_MUTATION)
    for i in range(GENE_SIZE):
        if i in idx:
            child[i]=np.random.randint(0,K,1)
def cross_and_mutation(gene_list):
    res=[]
    for father in gene_list:
        child=father
        if np.random.rand()<CROSS_RATE:
            mother=gene_list[np.random.randint(POPULATION_SIZE)]
            cross(child,mother)
        if np.random.rand()<MUTATION_RATE:
            mutation(child)
        res.append(child)
    return res
def init_population():
    res = []
    for i in range(POPULATION_SIZE):
        gene=np.zeros(GENE_SIZE)
        for j in range(GENE_SIZE):
            # gene[j] = np.random.randint(0,K,1)
            gene[j]=0
        res.append(gene)
    return res
# def get_best_gene(gene_list):
#     res=0
#     max_temp=0
#     for i in range(len(gene_list)):
#         score=get_score(gene_list[i])
#         if score>max_temp:
#             max_temp=score
#             res=i
#     return gene_list[res]
def print_pic(list):
    x=np.arange(len(list))
    plt.plot(x,list)
    plt.show()

def run(env):
    fit_list=[]
    gene_list = np.array(init_population())
    for i in range(EPOCH):
        cross_and_mutation(gene_list)
        fitness = get_fitness(gene_list,env)
        score=np.max(fitness)
        fit_list.append(score)
        print(score)
        gene_list=select(gene_list,fitness)

    print_pic(fit_list)


if __name__ == "__main__":
    env=network.ENV('cur',0)
    env.load_task()
    dqn = DQN(n_actions=network.K, n_features=network.SLOT_NUM,
              node_size=14,
              learning_rate=0.00005,
              reward_decay=1,
              e_greedy=0.95,
              replace_target_iter=20,
              memory_size=10,
              batch_size=5,
              e_greedy_increment=None
              )
    dff=eval_DFF(env,dqn)
    print('dff得分是'+str(dff))
    run(env)
